import numpy as np

class ConvergenceChecker:
    def __init__(self, tolerance=1e-6):
        self.tolerance = tolerance
        self.mass_balance_history = []
        self.energy_balance_history = []
        
    def check_mass_conservation(self, subchannel):
        """检查质量守恒"""
        mass_in = 0
        mass_out = 0
        
        # 计算入口和出口的质量流量
        for j in range(subchannel.n_channels):
            # 入口质量流量
            rho_l = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[0,j])[1]
            rho_v = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[0,j])[2]
            alpha = subchannel.void_fraction[0,j]
            
            mass_in += (1-alpha) * rho_l * subchannel.vel_liquid[0,j] + \
                      alpha * rho_v * subchannel.vel_vapor[0,j]
            
            # 出口质量流量
            rho_l = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[-1,j])[1]
            rho_v = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[-1,j])[2]
            alpha = subchannel.void_fraction[-1,j]
            
            mass_out += (1-alpha) * rho_l * subchannel.vel_liquid[-1,j] + \
                       alpha * rho_v * subchannel.vel_vapor[-1,j]
        
        mass_balance = abs(mass_in - mass_out) / mass_in
        self.mass_balance_history.append(mass_balance)
        
        return mass_balance < self.tolerance
        
    def check_energy_conservation(self, subchannel, heat_input):
        """检查能量守恒"""
        energy_in = 0
        energy_out = 0
        
        # 计算入口和出口的能量
        for j in range(subchannel.n_channels):
            # 入口能量
            h_l = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[0,j])[3]
            h_v = h_l + subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[0,j])[3]
            alpha = subchannel.void_fraction[0,j]
            
            energy_in += (1-alpha) * h_l + alpha * h_v
            
            # 出口能量
            h_l = subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[-1,j])[3]
            h_v = h_l + subchannel.fluid_prop.get_saturation_properties(subchannel.pressure[-1,j])[3]
            alpha = subchannel.void_fraction[-1,j]
            
            energy_out += (1-alpha) * h_l + alpha * h_v
            
        energy_balance = abs(energy_out - energy_in - heat_input) / heat_input
        self.energy_balance_history.append(energy_balance)
        
        return energy_balance < self.tolerance
        
    def generate_convergence_report(self):
        """生成收敛性报告"""
        report = {
            'mass_balance': {
                'final': self.mass_balance_history[-1],
                'history': self.mass_balance_history
            },
            'energy_balance': {
                'final': self.energy_balance_history[-1],
                'history': self.energy_balance_history
            }
        }
        return report 